// Copyright (c) 2003-2004 Brian Wellington (bwelling@xbill.org)
// Parts of this are derived from lib/dns/xfrin.c from BIND 9; its copyright
// notice follows.

/*
 * Copyright (C) 1999-2001  Internet Software Consortium.
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND INTERNET SOFTWARE CONSORTIUM
 * DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL
 * INTERNET SOFTWARE CONSORTIUM BE LIABLE FOR ANY SPECIAL, DIRECT,
 * INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING
 * FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
 * NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
 * WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

package org.xbill.DNS;

import java.io.*;
import java.net.*;
import java.util.*;

/**
 * An incoming DNS Zone Transfer. To use this class, first initialize an object,
 * then call the run() method. If run() doesn't throw an exception the result
 * will either be an IXFR-style response, an AXFR-style response, or an
 * indication that the zone is up to date.
 * 
 * @author Brian Wellington
 */

public class ZoneTransferIn {

	private static final int INITIALSOA = 0;
	private static final int FIRSTDATA = 1;
	private static final int IXFR_DELSOA = 2;
	private static final int IXFR_DEL = 3;
	private static final int IXFR_ADDSOA = 4;
	private static final int IXFR_ADD = 5;
	private static final int AXFR = 6;
	private static final int END = 7;

	private Name zname;
	private int qtype;
	private int dclass;
	private long ixfr_serial;
	private boolean want_fallback;
	private ZoneTransferHandler handler;

	private SocketAddress localAddress;
	private SocketAddress address;
	private TCPClient client;
	private TSIG tsig;
	private TSIG.StreamVerifier verifier;
	private long timeout = 900 * 1000;

	private int state;
	private long end_serial;
	private long current_serial;
	private Record initialsoa;

	private int rtype;

	public static class Delta {
		/**
		 * All changes between two versions of a zone in an IXFR response.
		 */

		/** The starting serial number of this delta. */
		public long start;

		/** The ending serial number of this delta. */
		public long end;

		/** A list of records added between the start and end versions */
		public List adds;

		/** A list of records deleted between the start and end versions */
		public List deletes;

		private Delta() {
			adds = new ArrayList();
			deletes = new ArrayList();
		}
	}

	public static interface ZoneTransferHandler {
		/**
		 * Handles a Zone Transfer.
		 */

		/**
		 * Called when an AXFR transfer begins.
		 */
		public void startAXFR() throws ZoneTransferException;

		/**
		 * Called when an IXFR transfer begins.
		 */
		public void startIXFR() throws ZoneTransferException;

		/**
		 * Called when a series of IXFR deletions begins.
		 * 
		 * @param soa
		 *            The starting SOA.
		 */
		public void startIXFRDeletes(Record soa) throws ZoneTransferException;

		/**
		 * Called when a series of IXFR adds begins.
		 * 
		 * @param soa
		 *            The starting SOA.
		 */
		public void startIXFRAdds(Record soa) throws ZoneTransferException;

		/**
		 * Called for each content record in an AXFR.
		 * 
		 * @param r
		 *            The DNS record.
		 */
		public void handleRecord(Record r) throws ZoneTransferException;
	};

	private static class BasicHandler implements ZoneTransferHandler {
		private List axfr;
		private List ixfr;

		public void startAXFR() {
			axfr = new ArrayList();
		}

		public void startIXFR() {
			ixfr = new ArrayList();
		}

		public void startIXFRDeletes(Record soa) {
			Delta delta = new Delta();
			delta.deletes.add(soa);
			delta.start = getSOASerial(soa);
			ixfr.add(delta);
		}

		public void startIXFRAdds(Record soa) {
			Delta delta = (Delta) ixfr.get(ixfr.size() - 1);
			delta.adds.add(soa);
			delta.end = getSOASerial(soa);
		}

		public void handleRecord(Record r) {
			List list;
			if (ixfr != null) {
				Delta delta = (Delta) ixfr.get(ixfr.size() - 1);
				if (delta.adds.size() > 0)
					list = delta.adds;
				else
					list = delta.deletes;
			} else
				list = axfr;
			list.add(r);
		}
	};

	private ZoneTransferIn() {
	}

	private ZoneTransferIn(Name zone, int xfrtype, long serial,
			boolean fallback, SocketAddress address, TSIG key) {
		this.address = address;
		this.tsig = key;
		if (zone.isAbsolute())
			zname = zone;
		else {
			try {
				zname = Name.concatenate(zone, Name.root);
			} catch (NameTooLongException e) {
				throw new IllegalArgumentException("ZoneTransferIn: "
						+ "name too long");
			}
		}
		qtype = xfrtype;
		dclass = DClass.IN;
		ixfr_serial = serial;
		want_fallback = fallback;
		state = INITIALSOA;
	}

	/**
	 * Instantiates a ZoneTransferIn object to do an AXFR (full zone transfer).
	 * 
	 * @param zone
	 *            The zone to transfer.
	 * @param address
	 *            The host/port from which to transfer the zone.
	 * @param key
	 *            The TSIG key used to authenticate the transfer, or null.
	 * @return The ZoneTransferIn object.
	 * @throws UnknownHostException
	 *             The host does not exist.
	 */
	public static ZoneTransferIn newAXFR(Name zone, SocketAddress address,
			TSIG key) {
		return new ZoneTransferIn(zone, Type.AXFR, 0, false, address, key);
	}

	/**
	 * Instantiates a ZoneTransferIn object to do an AXFR (full zone transfer).
	 * 
	 * @param zone
	 *            The zone to transfer.
	 * @param host
	 *            The host from which to transfer the zone.
	 * @param port
	 *            The port to connect to on the server, or 0 for the default.
	 * @param key
	 *            The TSIG key used to authenticate the transfer, or null.
	 * @return The ZoneTransferIn object.
	 * @throws UnknownHostException
	 *             The host does not exist.
	 */
	public static ZoneTransferIn newAXFR(Name zone, String host, int port,
			TSIG key) throws UnknownHostException {
		if (port == 0)
			port = SimpleResolver.DEFAULT_PORT;
		return newAXFR(zone, new InetSocketAddress(host, port), key);
	}

	/**
	 * Instantiates a ZoneTransferIn object to do an AXFR (full zone transfer).
	 * 
	 * @param zone
	 *            The zone to transfer.
	 * @param host
	 *            The host from which to transfer the zone.
	 * @param key
	 *            The TSIG key used to authenticate the transfer, or null.
	 * @return The ZoneTransferIn object.
	 * @throws UnknownHostException
	 *             The host does not exist.
	 */
	public static ZoneTransferIn newAXFR(Name zone, String host, TSIG key)
			throws UnknownHostException {
		return newAXFR(zone, host, 0, key);
	}

	/**
	 * Instantiates a ZoneTransferIn object to do an IXFR (incremental zone
	 * transfer).
	 * 
	 * @param zone
	 *            The zone to transfer.
	 * @param serial
	 *            The existing serial number.
	 * @param fallback
	 *            If true, fall back to AXFR if IXFR is not supported.
	 * @param address
	 *            The host/port from which to transfer the zone.
	 * @param key
	 *            The TSIG key used to authenticate the transfer, or null.
	 * @return The ZoneTransferIn object.
	 * @throws UnknownHostException
	 *             The host does not exist.
	 */
	public static ZoneTransferIn newIXFR(Name zone, long serial,
			boolean fallback, SocketAddress address, TSIG key) {
		return new ZoneTransferIn(zone, Type.IXFR, serial, fallback, address,
				key);
	}

	/**
	 * Instantiates a ZoneTransferIn object to do an IXFR (incremental zone
	 * transfer).
	 * 
	 * @param zone
	 *            The zone to transfer.
	 * @param serial
	 *            The existing serial number.
	 * @param fallback
	 *            If true, fall back to AXFR if IXFR is not supported.
	 * @param host
	 *            The host from which to transfer the zone.
	 * @param port
	 *            The port to connect to on the server, or 0 for the default.
	 * @param key
	 *            The TSIG key used to authenticate the transfer, or null.
	 * @return The ZoneTransferIn object.
	 * @throws UnknownHostException
	 *             The host does not exist.
	 */
	public static ZoneTransferIn newIXFR(Name zone, long serial,
			boolean fallback, String host, int port, TSIG key)
			throws UnknownHostException {
		if (port == 0)
			port = SimpleResolver.DEFAULT_PORT;
		return newIXFR(zone, serial, fallback,
				new InetSocketAddress(host, port), key);
	}

	/**
	 * Instantiates a ZoneTransferIn object to do an IXFR (incremental zone
	 * transfer).
	 * 
	 * @param zone
	 *            The zone to transfer.
	 * @param serial
	 *            The existing serial number.
	 * @param fallback
	 *            If true, fall back to AXFR if IXFR is not supported.
	 * @param host
	 *            The host from which to transfer the zone.
	 * @param key
	 *            The TSIG key used to authenticate the transfer, or null.
	 * @return The ZoneTransferIn object.
	 * @throws UnknownHostException
	 *             The host does not exist.
	 */
	public static ZoneTransferIn newIXFR(Name zone, long serial,
			boolean fallback, String host, TSIG key)
			throws UnknownHostException {
		return newIXFR(zone, serial, fallback, host, 0, key);
	}

	/**
	 * Gets the name of the zone being transferred.
	 */
	public Name getName() {
		return zname;
	}

	/**
	 * Gets the type of zone transfer (either AXFR or IXFR).
	 */
	public int getType() {
		return qtype;
	}

	/**
	 * Sets a timeout on this zone transfer. The default is 900 seconds (15
	 * minutes).
	 * 
	 * @param secs
	 *            The maximum amount of time that this zone transfer can take.
	 */
	public void setTimeout(int secs) {
		if (secs < 0)
			throw new IllegalArgumentException("timeout cannot be "
					+ "negative");
		timeout = 1000L * secs;
	}

	/**
	 * Sets an alternate DNS class for this zone transfer.
	 * 
	 * @param dclass
	 *            The class to use instead of class IN.
	 */
	public void setDClass(int dclass) {
		DClass.check(dclass);
		this.dclass = dclass;
	}

	/**
	 * Sets the local address to bind to when sending messages.
	 * 
	 * @param addr
	 *            The local address to send messages from.
	 */
	public void setLocalAddress(SocketAddress addr) {
		this.localAddress = addr;
	}

	private void openConnection() throws IOException {
		long endTime = System.currentTimeMillis() + timeout;
		client = new TCPClient(endTime);
		if (localAddress != null)
			client.bind(localAddress);
		client.connect(address);
	}

	private void sendQuery() throws IOException {
		Record question = Record.newRecord(zname, qtype, dclass);

		Message query = new Message();
		query.getHeader().setOpcode(Opcode.QUERY);
		query.addRecord(question, Section.QUESTION);
		if (qtype == Type.IXFR) {
			Record soa = new SOARecord(zname, dclass, 0, Name.root, Name.root,
					ixfr_serial, 0, 0, 0, 0);
			query.addRecord(soa, Section.AUTHORITY);
		}
		if (tsig != null) {
			tsig.apply(query, null);
			verifier = new TSIG.StreamVerifier(tsig, query.getTSIG());
		}
		byte[] out = query.toWire(Message.MAXLENGTH);
		client.send(out);
	}

	private static long getSOASerial(Record rec) {
		SOARecord soa = (SOARecord) rec;
		return soa.getSerial();
	}

	private void logxfr(String s) {
		if (Options.check("verbose"))
			System.out.println(zname + ": " + s);
	}

	private void fail(String s) throws ZoneTransferException {
		throw new ZoneTransferException(s);
	}

	private void fallback() throws ZoneTransferException {
		if (!want_fallback)
			fail("server doesn't support IXFR");

		logxfr("falling back to AXFR");
		qtype = Type.AXFR;
		state = INITIALSOA;
	}

	private void parseRR(Record rec) throws ZoneTransferException {
		int type = rec.getType();
		Delta delta;

		switch (state) {
		case INITIALSOA:
			if (type != Type.SOA)
				fail("missing initial SOA");
			initialsoa = rec;
			// Remember the serial number in the initial SOA; we need it
			// to recognize the end of an IXFR.
			end_serial = getSOASerial(rec);
			if (qtype == Type.IXFR
					&& Serial.compare(end_serial, ixfr_serial) <= 0) {
				logxfr("up to date");
				state = END;
				break;
			}
			state = FIRSTDATA;
			break;

		case FIRSTDATA:
			// If the transfer begins with 1 SOA, it's an AXFR.
			// If it begins with 2 SOAs, it's an IXFR.
			if (qtype == Type.IXFR && type == Type.SOA
					&& getSOASerial(rec) == ixfr_serial) {
				rtype = Type.IXFR;
				handler.startIXFR();
				logxfr("got incremental response");
				state = IXFR_DELSOA;
			} else {
				rtype = Type.AXFR;
				handler.startAXFR();
				handler.handleRecord(initialsoa);
				logxfr("got nonincremental response");
				state = AXFR;
			}
			parseRR(rec); // Restart...
			return;

		case IXFR_DELSOA:
			handler.startIXFRDeletes(rec);
			state = IXFR_DEL;
			break;

		case IXFR_DEL:
			if (type == Type.SOA) {
				current_serial = getSOASerial(rec);
				state = IXFR_ADDSOA;
				parseRR(rec); // Restart...
				return;
			}
			handler.handleRecord(rec);
			break;

		case IXFR_ADDSOA:
			handler.startIXFRAdds(rec);
			state = IXFR_ADD;
			break;

		case IXFR_ADD:
			if (type == Type.SOA) {
				long soa_serial = getSOASerial(rec);
				if (soa_serial == end_serial) {
					state = END;
					break;
				} else if (soa_serial != current_serial) {
					fail("IXFR out of sync: expected serial " + current_serial
							+ " , got " + soa_serial);
				} else {
					state = IXFR_DELSOA;
					parseRR(rec); // Restart...
					return;
				}
			}
			handler.handleRecord(rec);
			break;

		case AXFR:
			// Old BINDs sent cross class A records for non IN classes.
			if (type == Type.A && rec.getDClass() != dclass)
				break;
			handler.handleRecord(rec);
			if (type == Type.SOA) {
				state = END;
			}
			break;

		case END:
			fail("extra data");
			break;

		default:
			fail("invalid state");
			break;
		}
	}

	private void closeConnection() {
		try {
			if (client != null)
				client.cleanup();
		} catch (IOException e) {
		}
	}

	private Message parseMessage(byte[] b) throws WireParseException {
		try {
			return new Message(b);
		} catch (IOException e) {
			if (e instanceof WireParseException)
				throw (WireParseException) e;
			throw new WireParseException("Error parsing message");
		}
	}

	private void doxfr() throws IOException, ZoneTransferException {
		sendQuery();
		while (state != END) {
			byte[] in = client.recv();
			Message response = parseMessage(in);
			if (response.getHeader().getRcode() == Rcode.NOERROR
					&& verifier != null) {
				TSIGRecord tsigrec = response.getTSIG();

				int error = verifier.verify(response, in);
				if (error != Rcode.NOERROR)
					fail("TSIG failure");
			}

			Record[] answers = response.getSectionArray(Section.ANSWER);

			if (state == INITIALSOA) {
				int rcode = response.getRcode();
				if (rcode != Rcode.NOERROR) {
					if (qtype == Type.IXFR && rcode == Rcode.NOTIMP) {
						fallback();
						doxfr();
						return;
					}
					fail(Rcode.string(rcode));
				}

				Record question = response.getQuestion();
				if (question != null && question.getType() != qtype) {
					fail("invalid question section");
				}

				if (answers.length == 0 && qtype == Type.IXFR) {
					fallback();
					doxfr();
					return;
				}
			}

			for (int i = 0; i < answers.length; i++) {
				parseRR(answers[i]);
			}

			if (state == END && verifier != null && !response.isVerified())
				fail("last message must be signed");
		}
	}

	/**
	 * Does the zone transfer.
	 * 
	 * @param handler
	 *            The callback object that handles the zone transfer data.
	 * @throws IOException
	 *             The zone transfer failed to due an IO problem.
	 * @throws ZoneTransferException
	 *             The zone transfer failed to due a problem with the zone
	 *             transfer itself.
	 */
	public void run(ZoneTransferHandler handler) throws IOException,
			ZoneTransferException {
		this.handler = handler;
		try {
			openConnection();
			doxfr();
		} finally {
			closeConnection();
		}
	}

	/**
	 * Does the zone transfer.
	 * 
	 * @return A list, which is either an AXFR-style response (List of Records),
	 *         and IXFR-style response (List of Deltas), or null, which
	 *         indicates that an IXFR was performed and the zone is up to date.
	 * @throws IOException
	 *             The zone transfer failed to due an IO problem.
	 * @throws ZoneTransferException
	 *             The zone transfer failed to due a problem with the zone
	 *             transfer itself.
	 */
	public List run() throws IOException, ZoneTransferException {
		BasicHandler handler = new BasicHandler();
		run(handler);
		if (handler.axfr != null)
			return handler.axfr;
		return handler.ixfr;
	}

	private BasicHandler getBasicHandler() throws IllegalArgumentException {
		if (handler instanceof BasicHandler)
			return (BasicHandler) handler;
		throw new IllegalArgumentException("ZoneTransferIn used callback "
				+ "interface");
	}

	/**
	 * Returns true if the response is an AXFR-style response (List of Records).
	 * This will be true if either an IXFR was performed, an IXFR was performed
	 * and the server provided a full zone transfer, or an IXFR failed and
	 * fallback to AXFR occurred.
	 */
	public boolean isAXFR() {
		return (rtype == Type.AXFR);
	}

	/**
	 * Gets the AXFR-style response.
	 * 
	 * @throws IllegalArgumentException
	 *             The transfer used the callback interface, so the response was
	 *             not stored.
	 */
	public List getAXFR() {
		BasicHandler handler = getBasicHandler();
		return handler.axfr;
	}

	/**
	 * Returns true if the response is an IXFR-style response (List of Deltas).
	 * This will be true only if an IXFR was performed and the server provided
	 * an incremental zone transfer.
	 */
	public boolean isIXFR() {
		return (rtype == Type.IXFR);
	}

	/**
	 * Gets the IXFR-style response.
	 * 
	 * @throws IllegalArgumentException
	 *             The transfer used the callback interface, so the response was
	 *             not stored.
	 */
	public List getIXFR() {
		BasicHandler handler = getBasicHandler();
		return handler.ixfr;
	}

	/**
	 * Returns true if the response indicates that the zone is up to date. This
	 * will be true only if an IXFR was performed.
	 * 
	 * @throws IllegalArgumentException
	 *             The transfer used the callback interface, so the response was
	 *             not stored.
	 */
	public boolean isCurrent() {
		BasicHandler handler = getBasicHandler();
		return (handler.axfr == null && handler.ixfr == null);
	}

}
